if (!file.exists("output")) {
  dir.create("output")
}

# ANOVA_regressions Performs logistic regression analysis of a full and reduced model and compares the fit of the models with ANOVA
# Dichotomized signature attributions are dependent variables and epidemiological factors are independent variables.
# Firth penalized logistic regression implemented if parameter or confounder has perfect or near perfect separation
# The function returns a data frame with ANOVA results

# data = data frame with all variables included in the model
# signatures = character vector with signature names
# parameters = character vector with parameter names
# confounders_full = character vector with parameter names included in the full model
# confounders_reduced = character vector with parameter names included in the reduced model

ANOVA_regressions <- function(data, signatures, parameters, confounders_full, confounders_reduced) {
  output <- data.frame()
  for (s in signatures) {
    for (p in parameters) {
      # check if requirements for Firth regression are met
      if (is.numeric(data[, p])) {
        next
      }
      firth_method <- F
      for (category in data[, p] %>% unique()) {
        category_counts <- data[data[, p] == category, ][s] %>%
          group_by(get(s)) %>%
          tally()
        if (length(category_counts$n) == 1 | category_counts$n[1] <= 1 | category_counts$n[2] <= 1) {
          print(paste("** Factor", p, "in signature", s, "has perfect or near-perfect separation for one or more levels. Using the penalised approach (Firth method)."))
          firth_method <- T
          break
        }
      }
      
      if (!firth_method) {
        for (c in confounders_full[confounders_full != "age"]) { ## age is a continuous variable so should not be taken into account for this
          for (category in data[, c] %>% unique()) {
            category_counts <- data[data[, c] == category, ][s] %>%
              group_by(get(s)) %>%
              tally()
            if (length(category_counts$n) == 1 | category_counts$n[1] <= 1 | category_counts$n[2] <= 1) {
              print(paste("** Covariate", c, "in signature", s, "has perfect or near-perfect separation for one or more levels. Using the penalised approach (Firth method)."))
              firth_method <- T
              break
            }
          }
          
          if (firth_method) break
        }
      }
      formula_reduced <- paste0(s, " ~ ", paste(c(parameters, confounders_reduced), collapse = " + "))
      formula_full <- paste0(s, " ~ ", paste(c(parameters, confounders_full), collapse = " + "))
      
      if (firth_method == F) {
        glm_reduced <- glm(formula_reduced, data = data, family = binomial("logit"))
        glm_full <- glm(formula_full, data = data, family = binomial("logit"))
        
        logLik_reduced <- logLik(glm_reduced)
        logLik_full <- logLik(glm_full)
      } else {
        firth_reduced <- logistf(formula = formula_reduced, data = data)
        firth_full <- logistf(formula = formula_full, data = data)
        
        logLik_reduced <- firth_reduced$loglik[1]
        logLik_full <- firth_full$loglik[1]
      }
      
      # Calculate the test statistic
      likelihood_ratio_stat <- -2 * (logLik_reduced - logLik_full)
      
      # Compare to a chi-squared distribution
      # The difference in degrees of freedom is 1 (model2 has one more parameter than model1)
      p_value <- pchisq(likelihood_ratio_stat, df = 1, lower.tail = FALSE)
      
      anova_summary <- data.frame(
        signature = rep(s, 2),
        model = c("reduced", "full"),
        model_function = c(formula_reduced, formula_full),
        log_likelihood = c(logLik_reduced, logLik_full),
        stat = rep(likelihood_ratio_stat, 2),
        p_val = rep(p_value, 2)
      )
      
      output <- rbind(output, anova_summary)
    }
  }
  return(output)
}

# load data
source("HNC_metadata_tidy.R")
source("data_handling.R")
source("risk_factor_regressions.R")

sbsCOSMIC <- read.csv("Input_signature_attributions/Input_SBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
dbsCOSMIC <- read.csv("Input_signature_attributions/Input_DBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
idCOSMIC <- read.csv("Input_signature_attributions/Input_ID_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")

additional_df <- list(sbsCOSMIC, dbsCOSMIC, idCOSMIC)

cosmic <- additional_df[[1]]
for (n in 2:length(additional_df)) {
  cosmic <- cosmic %>% left_join(additional_df[[n]], by = "donor_id")
}

# merge with cosmic sigantures and tidy epi variables
data <- data %>%
  left_join(cosmic, by = "donor_id") %>%
  mutate(
    country = factor(country, levels = c("Brazil", "Colombia", "Argentina", "Greece", "Italy", "Czech Republic", "Slovakia", "Romania")),
    argentina = ifelse(country == "Argentina", "yes", "no"),
    greece = ifelse(country == "Greece", "yes", "no"),
    italy = ifelse(country == "Italy", "yes", "no"),
    brazil = ifelse(country == "Brazil", "yes", "no"),
    czechRepublic = ifelse(country == "Czech Republic", "yes", "no"),
    romania = ifelse(country == "Romania", "yes", "no"),
    subsite = factor(subsite, levels = c("Larynx", "OC", "OPC", "Hypopharynx")),
    age_group = factor(age_group, levels = c("55-65", "0-45", "45-55", "65-75", "75+")),
    age = as.numeric(scale(age_diag))
  )

# Select signatures present in >n% of samples
# Dichotomize the signatures (Presence (signatures>0)="1",absence (signatures==0)="0")
# If signature is present in more than 75% cases, dichotomize by the median
data <- dichotomizeSigs(metadata = data, sigsn = cosmic, sigcutoff = 2)

# Dichotomized signatures to analyze
# signatures <- grep("^(SBS|DBS|ID).*_cat$", names(data), value = T)
signatures <- c("SBS16_cat", "DBS4_cat", "ID11_cat")

parameters <- "alcohol_ever"

confounders_full <- c("tobacco_ever", "age", "sex", "region", "subsite")
confounders_reduced <- c("age", "sex", "region", "subsite")

ANOVA <- ANOVA_regressions(data, signatures, parameters, confounders_full, confounders_reduced)

write.csv(ANOVA, "output/Supplementary_Note_Table_2.csv", row.names = F)
